#include <eigen3/Eigen/Dense>
#include <vector>
#include <iostream>
#define BACKGROUND_COLOR Vector3f(0, 0, 0)
#define MAX_DEPTH 3

using namespace Eigen;
using namespace std;

struct Ray
{
    Vector3f origin;
    Vector3f direction;
};
struct Light
{
    Vector3f position;
    Vector3f intensity;
};

struct Intersection
{
    bool hit;                // 是否命中
    float t;                 // 光线参数 t
    Vector3f point;          // 相交点
    Vector3f normal;         // 法线
    Vector3f color;          // 命中物体颜色
    float kd = 1.0f;         // 漫反射系数
    float ks = 0.5f;         // 镜面反射系数
    float shininess = 32.0f; // 高光指数
    float kr = 0.2f;         // 反射系数
};

class Object
{
public:
    virtual Intersection intersect(Ray ray) = 0;
};

class Plane : public Object
{
public:
    Vector3f point;
    Vector3f normal;
    Vector3f color;
    Plane(Vector3f p, Vector3f n, Vector3f c) : point(p), normal(n), color(c) {};
    Intersection intersect(Ray ray)
    {
        Intersection result;
        result.hit = false;
        float denom = normal.dot(ray.direction);
        if (fabs(denom) > 1e-6)
        {
            float t = (point - ray.origin).dot(normal) / denom;
            if (t >= 0)
            {
                result.hit = true;
                result.t = t;
                result.point = ray.origin + t * ray.direction;
                result.normal = normal.normalized();
                result.color = color;
            }
        }
        return result;
    }
};

class Sphere:public Object
{
public:
    Vector3f center;
    float radius;
    Vector3f color;
    Sphere(Vector3f c, float r, Vector3f col) : center(c), radius(r), color(col) {};
    Intersection intersect(Ray ray)
    {
        Intersection result;
        result.hit = false;
        Vector3f oc = ray.origin - center;
        float b = 2 * oc.dot(ray.direction);
        float c = oc.dot(oc) - radius * radius;
        float delta = b * b - 4 * c;
        if (delta >= 0)
        {
            float t1 = (-b - sqrt(delta)) / 2;
            float t2 = (-b + sqrt(delta)) / 2;
            float t = -1;
            if (t1 > 1e-4)
                t = t1;
            else if (t2 > 1e-4)
                t = t2;
            if (t > 0)
            {
                result.hit = true;
                result.t = t;
                result.point = ray.origin + t * ray.direction;
                result.normal = (result.point - center).normalized();
                result.color = color;
            }
        }
        return result;
    }
};

class Scene
{
public:
    vector<Object*> objects;
    vector<Light> lights;
    Vector3f ambientLight = Vector3f(0.2, 0.2, 0.2); // 环境光强度

    bool isInShadow(const Vector3f &point, const Light &light)
    {
        Vector3f lightDir = (light.position - point).normalized();
        Ray shadowRay;
        shadowRay.origin = point + 1e-4 * lightDir; // 避免自阴影
        shadowRay.direction = lightDir;
        for (auto &obj : objects)
        {
            Intersection result = obj->intersect(shadowRay);
            if (result.hit)
                return true;
        }
        return false;
    }

    Vector3f trace(Ray ray, int depth = 0)
    {
        if (depth > MAX_DEPTH)
        {
            return BACKGROUND_COLOR;
        }
        float min_t = 1e10;
        Vector3f color = BACKGROUND_COLOR;
        Intersection closestIntersection;
        closestIntersection.hit = false;

        // 找最近交点
        for (int i = 0; i < objects.size(); i++)
        {
            Intersection result = objects[i]->intersect(ray);
            if (result.hit && result.t < min_t)
            {
                min_t = result.t;
                closestIntersection = result;
            }
        }
        if (!closestIntersection.hit)
        {
            return BACKGROUND_COLOR;
        }
        else
        {
            Vector3f finalColor = closestIntersection.color.cwiseProduct(ambientLight); // 环境光
            Vector3f diffuseSum(0, 0, 0);
            Vector3f specularSum(0, 0, 0);
            for (int i = 0; i < lights.size(); i++)
            {
                // 阴影检测
                if (isInShadow(closestIntersection.point, lights[i]))
                    continue;
                // 漫反射
                Vector3f lightDir = (lights[i].position - closestIntersection.point).normalized();
                float diff = max(0.0f, closestIntersection.normal.dot(lightDir));
                diffuseSum += diff * lights[i].intensity;
                // 镜面高光
                Vector3f reflectDir = (2 * closestIntersection.normal.dot(lightDir) * closestIntersection.normal - lightDir).normalized();
                float spec = pow(max(0.0f, reflectDir.dot(-ray.direction)), closestIntersection.shininess);
                specularSum += lights[i].intensity * spec;
            }
            // 乘上材质系数
            Vector3f diffuseColor = closestIntersection.kd * diffuseSum.cwiseProduct(closestIntersection.color);
            Vector3f specularColor = closestIntersection.ks * specularSum;
            finalColor += diffuseColor + specularColor;

            // 反射
            if (closestIntersection.kr > 0)
            {
                Vector3f reflectDir = (ray.direction.normalized() - 2 * closestIntersection.normal.normalized().dot(ray.direction) * closestIntersection.normal.normalized()).normalized();
                Ray reflectRay;
                reflectRay.origin = closestIntersection.point + 1e-4 * (reflectDir);
                reflectRay.direction = reflectDir;
                Vector3f reflectColor = trace(reflectRay, depth + 1);
                finalColor = (1.0f - closestIntersection.kr) * finalColor + reflectColor * closestIntersection.kr;
            }

            finalColor = finalColor.cwiseMin(Vector3f(1, 1, 1));
            return finalColor;
        }
    }

    void addObject(Object *obj)
    {
        objects.push_back(obj);
    }
    void addLight(Light light)
    {
        lights.push_back(light);
    }
};

class Camera
{
public:
    Matrix3f intrinsic;
    Matrix4f extrinsic;
    Ray generateRay(float x, float y)
    {
        Vector3f pixel_homogeneous(x, y, 1);
        Vector3f direction = intrinsic.inverse() * pixel_homogeneous;
        direction.normalize();
        Vector4f direction_homogeneous(direction.x(), direction.y(), direction.z(), 0);
        Vector4f origin_homogeneous = extrinsic * Vector4f(0, 0, 0, 1);
        Vector4f dir_world = extrinsic * direction_homogeneous;
        Ray ray;
        ray.origin = origin_homogeneous.head<3>();
        ray.direction = dir_world.head<3>();
        ray.direction.normalize();
        return ray;
    }
};

int main()
{
    Scene scene;
    Camera camera;
    int width = 800;
    int height = 800;
    camera.intrinsic << width, 0, width / 2,
        0, height, height / 2,
        0, 0, 1;
    camera.extrinsic << 1, 0, 0, 0,
        0, 1, 0, 0,
        0, 0, 1, 0,
        0, 0, 0, 1;
    scene.addObject(new Sphere{Vector3f(0, 0, 3), 1, Vector3f(1, 0, 0)});
    scene.addObject(new Sphere{Vector3f(2, 0, 6), 1, Vector3f(0, 1, 0)});
    scene.addObject(new Sphere{Vector3f(-2, 0, 6), 1, Vector3f(0, 0, 1)});
    scene.addObject(new Plane{Vector3f(0, 1, 0), Vector3f(0, -1, 0), Vector3f(1, 1, 0)});
    scene.addLight(Light{Vector3f(5, -1, 1), Vector3f(1, 1, 1)});
    vector<Vector3f> framebuffer(width * height);
    for (int y = 0; y < height; y++)
    {
        for (int x = 0; x < width; x++)
        {
            Ray ray = camera.generateRay(x, y);
            framebuffer[y * width + x] = scene.trace(ray);
        }
    }
    FILE *f = fopen("output.ppm", "w");
    fprintf(f, "P3\n%d %d\n255\n", width, height);
    for (int i = 0; i < width * height; i++)
    {
        fprintf(f, "%d %d %d ", int(framebuffer[i].x() * 255), int(framebuffer[i].y() * 255), int(framebuffer[i].z() * 255));
    }
    fclose(f);
    return 0;
}
